import itertools
from Create_more_metrics import create_new_metrics
from evolve import Evolution, query_fitness, query_fitness_for_indi
from sklearn import svm
from population import Population
import numpy as np
import copy
from get_data_from_101 import delete_margin
from utils import operations2onehot, GP_log, population_log, NULL, write_best_individual
from nasbench import api
from individual import Individual
import time
import random
import pickle
from sklearn.neural_network import MLPClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.ensemble import RandomForestRegressor
# transform genotype population to phenotype, and pad the matrix and the op_list to max length
def genotype2phenotype(pops: Population) -> Population:
    genotype_population = pops
    phenotype_population = Population(0)
    for indi in genotype_population.pops:
        matrix = indi.indi['matrix']
        op_list = indi.indi['op_list']
        model_spec = api.ModelSpec(matrix, op_list)
        pruned_matrix = model_spec.matrix
        pruned_op_list = model_spec.ops
        # start to padding zeros to 7*7 matrix and 7 op_list
        len_operations = len(pruned_op_list)
        assert len_operations == len(pruned_matrix)
        padding_matrix = copy.deepcopy(pruned_matrix)
        if len_operations != 7:
            for j in range(len_operations, 7):
                pruned_op_list.insert(j - 1, NULL)

            padding_matrix = np.insert(pruned_matrix, len_operations - 1,
                                       np.zeros([7 - len_operations, len_operations]), axis=0)
            padding_matrix = np.insert(padding_matrix, [len_operations - 1], np.zeros([7, 7 - len_operations]), axis=1)
        phenotype_individual = Individual()
        phenotype_individual.create_an_individual(padding_matrix, pruned_op_list)
        phenotype_individual.mean_acc = indi.mean_acc
        phenotype_individual.mean_training_time = indi.mean_training_time
        phenotype_population.add_individual(phenotype_individual)
    return phenotype_population


def get_input_X(pops: Population):
    X = []
    for indi in pops.pops:
        input_metrix = []
        matrix = indi.indi['matrix']
        matrix = delete_margin(np.array(matrix))
        flattened_adjacent = []
        matrix_size = len(matrix)
        # get upper triangular data in matrix and flat it
        # Noting that this doesn't contains the elements as main diagonal
        for i in range(matrix_size):
            flattened_adjacent.extend(matrix[i][i:])
        input_metrix.extend(flattened_adjacent)

        op_list = indi.indi['op_list']
        op_list = operations2onehot(op_list[1: -1])
        input_metrix.extend(op_list)

        X.append(input_metrix)
    return X


# Matrix and Op_list to vector
def MO2vector(matrix, op_list):
    input_metrix = []
    X=[]
    matrix = delete_margin(matrix)
    arraymatrix=np.array(matrix)
    flattened_adjacent = []
    matrix_size = len(matrix)
    # get upper triangular data in matrix and flat it
    # Noting that this doesn't contains the elements as main diagonal
    for i in range(matrix_size):
        flattened_adjacent.extend(matrix[i][i:])
    input_metrix.extend(flattened_adjacent)

    op_list = operations2onehot(op_list)
    input_metrix.extend(op_list)

    return input_metrix

def get_pair(X):
    assert len(X) > 0
    num = len(X)
    data=[]
    data_x = list(itertools.permutations(X,r=3))
    data = np.array(data_x)
    return data.reshape(-1,123)
# Get cpu, gpu or mps device for training.

# this is for genotype surrogate evolution
class SVM_evolution(Evolution):
    def __init__(self,  m_prob=0.2, m_num_matrix=1, m_num_op_list=1, x_prob=0.9,
                 population_size=100):
        super(SVM_evolution, self).__init__(m_prob, m_num_matrix, m_num_op_list, x_prob, population_size)
        self.archive = Population(0)
    def init_MLP_model(self):
        self.MLPmodel= MLPClassifier(solver='adam',
                                     activation='relu',
                                     hidden_layer_sizes=(246,123,32),
                                    batch_size=128,
                                     warm_start=True,
                                     verbose=True)
    def init_SVM_model(self):
        self.SVMmodel= svm.SVC()
    def init_RF_model(self):
        self.RFmodel = RandomForestRegressor(n_estimators=230)

    def fit_RF_model(self,argument=True):
        X = []
        y = []
        sum=0
        for indi in self.archive.pops:
            matrix = indi.indi['matrix']
            op_list = indi.indi['op_list']
            if argument:
                sum+=1
                more_metrics = create_new_metrics(matrix, op_list[1: -1], select_upper_tri=False, max_num=120)
                for same_metric in more_metrics:
                    adjacent_matrix, module_integers = same_metric['module_adjacency'], same_metric['module_integers']
                    input_metrix = MO2vector(adjacent_matrix, module_integers)
                    X.append(input_metrix)
                    y.append(indi.mean_acc)
            else:
                input_metrix = MO2vector(matrix, op_list[1: -1])
                X.append(input_metrix)
                y.append(indi.mean_acc)
        print("优秀个体数：",sum)
        self.RFmodel.fit(X, y)
    def predict_by_RF(self, pred_pops: Population, phenotype=False):
        if phenotype:
            pred_pops_new = genotype2phenotype(pred_pops)
            X = get_input_X(pred_pops_new)
        else:
            X = get_input_X(pred_pops)
        pred_y = self.RFmodel.predict(X)
        for i, indi in enumerate(pred_pops.pops):
            pred_pops.pops[i].mean_acc = pred_y[i]
            # using the GP model, don't need to training
            pred_pops.pops[i].mean_training_time = 0
    def fit_SVM_model(self, data_augmentation=True):
        Y = []
        label=[]
        inputx = get_input_X(self.archive)
        for indi in self.archive.pops:
            label.append(indi.mean_acc)
        X1=list(itertools.permutations(inputx,r=3))
        Y1 = list(itertools.permutations(label,r=3))
        X2 = np.array(X1)
        Y2 =np.array(Y1)
        assert len(X2)== len(Y2)
        for i,y in enumerate(Y2):
            if y[0] >= y[1] and y[0] >= y[2]:
                Y.append(1)
            else:
                Y.append(0)
        X = X2.reshape(-1, 123)
        self.MLPmodel.fit(X, Y)
        print("Done\n")

        # f2 = open('svm300.model', 'rb')
        # s2 = f2.read()
        # model1 = pickle.loads(s2)

    def predict_by_MLP(self, pred_pops: Population, phenotype=False,population_size = 100):
        if phenotype:
            pred_pops_new = genotype2phenotype(pred_pops)
            X = get_input_X(pred_pops_new)
        else:
            X = get_input_X(pred_pops)
        test_X = get_pair(X)
        pred_y =  self.MLPmodel.predict(test_X)
        # pred_y = model1.predict(test_X)
        y=np.array_split(pred_y,population_size)
        for i, j in enumerate(y):
            k = np.sum(j)
            pred_pops.pops[i].win(k)
    def add_fit_random(self, add_num):
        X = []
        Y = []
        label =[]
        arr=np.arange(population_size)
        np.random.shuffle(arr)
        z = Population(0)
        for i in range(add_num):
            z.add_individual(self.archive.get_individual_at(arr[i]))
        inputx = get_input_X(z)
        for indi in z.pops:
            label.append(indi.mean_acc)
        X1 = list(itertools.permutations(inputx, r=3))
        Y1 = list(itertools.permutations(label, r=3))
        X = np.array(X1)
        Y2 = np.array(Y1)
        assert len(X) == len(Y2)
        for i in range(len(X)):
            y = np.array(Y2[i])
            if y[0] >= y[1] and y[0] >= y[2]:
                Y.append(1)
            else:
                Y.append(0)
        X = X.reshape(-1, 123)
        try:
            self.MLPmodel.fit(X, Y)
            print("Done\n")
        except:
            print('ValueError: The number of classes has to be greater than one; got 1 class')
    def add_fit_accordranker(self, add_num):
        X = []
        Y = []
        label =[]
        a=self.archive.get_sorted_index_order_by_win()
        b=Population(0)
        for i in range(add_num):
            b.add_individual(self.archive.get_individual_at(a[i]))

        inputx = get_input_X(b)
        for indi in b.pops:
            label.append(indi.mean_acc)
        X1 = list(itertools.permutations(inputx, r=3))
        Y1 = list(itertools.permutations(label, r=3))
        X = np.array(X1)
        Y2 = np.array(Y1)
        assert len(X) == len(Y2)
        for i in range(len(X)):
            y = np.array(Y2[i])
            if y[0] >= y[1] and y[0] >= y[2]:
                Y.append(1)
            else:
                Y.append(0)
        X = X.reshape(-1, 123)
        try:
            self.MLPmodel.fit(X, Y)
            print("Done\n")
        except:
          print('ValueError: The number of classes has to be greater than one; got 1 class')


if __name__ == '__main__':
    for i in range(10):
        # time_begin = time.time()
        m_prob = 0.2
        m_num_matrix = 1
        m_num_op_list = 1
        x_prob = 0.8
        population_size = 100  # 必须是偶数
        num_generation = 20
        num_resample = 0
        archive_num = 210
        add_num = 20
        # phenotype=True means using phenotype to predict and fitting the GPmodel with phenotype
        # phenotype=False means using genotype to predict and fitting the GPmodel with genotype
        phenotype = True
        surrogate = True

        total_training_time = 0
        final_one_acc = []

        SVM_Evolution = SVM_evolution(m_prob, m_num_matrix, m_num_op_list, x_prob, population_size)
        # SVM_Evolution.initialize_popualtion(type='GP')
        SVM_Evolution.initialize_popualtion()
        SVM_Evolution.init_MLP_model()
        SVM_Evolution.init_SVM_model()
        gen_no = 0
        # query_fitness(gen_no, RF_Evolution.pops)
        # total_training_time += RF_Evolution.pops.calculate_population_training_time()

        if surrogate:
            fit_archive = Population(archive_num)
            query_fitness(-1, fit_archive)
            if not phenotype:
                SVM_Evolution.archive = copy.deepcopy(fit_archive)
            else:  # if phenotype:
                SVM_Evolution.archive = genotype2phenotype(fit_archive)
                print("-------------MLPfit---------------")
                SVM_Evolution.fit_SVM_model()
        # f2 = open('svm.model', 'rb')
        # s2 = f2.read()
        # model1 = pickle.loads(s2)

        if not surrogate:
            query_fitness(gen_no, SVM_Evolution.pops)
        else:
            print("-------------pred---------------")
            SVM_Evolution.predict_by_MLP(SVM_Evolution.pops, phenotype=phenotype,population_size=population_size)

        while True:
            gen_no += 1
            if gen_no > num_generation:
                break
            offspring = SVM_Evolution.recombinate(population_size)    # 从SVM_Evo.pops中搞出offspring
            # print('OFF',offspring)
            #print(offspring.get_pop_size())
            #print(offspring)
            #SVM_Evolution.pops.merge_populations(offspring.pops)
            #SVM_Evolution.predict_by_MLP(SVM_Evolution.pops, phenotype=phenotype)
            SVM_Evolution.environmental_selection(gen_no, offspring)

            SVM_Evolution.predict_by_MLP(SVM_Evolution.pops, phenotype=phenotype, population_size=population_size)
            query_fitness(gen_no, SVM_Evolution.pops)
            #SVM_Evolution.predict_by_MLP(SVM_Evolution.pops, phenotype=phenotype, population_size=population_size)
            #SVM_Evolution.pops.pops.sort(key=lambda x: x.win_time, reverse=True)  # 降序

            # # create the save_pops, because the self.pops may contains some individuals that don't have the real acc,
            # # actually, they contains the predicted acc
            # save_pops = copy.deepcopy(self.pops)
            # query_fitness(gen_no, save_pops)
            # population_log(gen_no, save_pops)
            population_log(gen_no, SVM_Evolution.pops)
            write_best_individual(gen_no, SVM_Evolution.pops)
            # 迭代训练
            add_fit_archive = SVM_Evolution.pops
            query_fitness(-1, add_fit_archive)
            SVM_Evolution.archive = genotype2phenotype(add_fit_archive)
            print("-------------fit_add---------------")
            if 10<gen_no<=15:
                SVM_Evolution.add_fit_random(add_num)
            elif 15<gen_no<=20:
                SVM_Evolution.add_fit_accordranker(add_num)
            # for the last generation
        last_resample_num = 4
        #sorted_win_index = SVM_Evolution.pops.get_sorted_index_order_by_win()
        final_population = Population(0)
        for i in range(4):
            final_population.add_individual(SVM_Evolution.pops.get_individual_at(i))
        #SVM_Evolution.predict_by_RF(last_population, phenotype=phenotype)
        #sorted_acc_index = last_population.get_sorted_index_order_by_acc()
        #final_population = Population(0)
        #for i in sorted_acc_index[:last_resample_num]:
            #final_population.add_individual(last_population.get_individual_at(i))
        gen_no = 'final'
        query_fitness(gen_no, final_population)
        population_log(gen_no, final_population)

        final_one_acc.append(final_population.pops[0].mean_acc)
        final_one_acc.append(final_population.pops[1].mean_acc)
        final_one_acc.append(final_population.pops[2].mean_acc)

        save_path = r'/data/public/zaj/try/GAon101/pop_log/total_training_time.txt'
        with open(save_path, 'w') as file:
            file.write('Total_training_time: ' + str(total_training_time) + '\n')
            file.write('Total_training_num: ' + str(population_size + num_generation * num_resample) + '\n')
            file.write(
                'm_prob: {}, m_num_matrix: {}, m_num_op_list: {}, x_prob: {}\n'.format(m_prob, m_num_matrix, m_num_op_list,
                                                                                       x_prob))
            file.write('RF_surrogate: True, num_resample: {}, phenotype: {}'.format(num_resample, phenotype))

        print('The ACC1 of the best individual found is: {}'.format(final_one_acc[0]))
        print('The ACC2 : {}'.format(final_one_acc[1]))
        print('The ACC3 : {}'.format(final_one_acc[2]))

        # time_end = time.time()
        # time = time_end - time_begin
        # print('time:', time)